#include "cuda_runtime.h"
#include "device_launch_parameters.h"

#include "sum_mean.hpp"
#include "stdio.h"
#include "stdlib.h"
#include "string.h"

static unsigned int nextPowerOfTwo(unsigned int x)
{
	--x;
	x |= x >> 1;
	x |= x >> 2;
	x |= x >> 4;
	x |= x >> 8;
	x |= x >> 16;
	return ++x;
}

__global__ void cuda_mean_kernel(uchar * devDataIn, float * devBuff, unsigned int length)
{
	int index = blockIdx.x * blockDim.x + threadIdx.x;

	if (index >= length)
	{
		return;
	}
	int halfBlockDim = (blockDim.x >> 1);
	extern __shared__ float sh_sum[];
	if (threadIdx.x < halfBlockDim)
	{
		uchar secVal = (index + halfBlockDim) < length ? devDataIn[index + halfBlockDim] : 0;
		sh_sum[threadIdx.x] = devDataIn[index] + secVal;
	}
	__syncthreads();

	for (int i = (blockDim.x >> 2); i > 0; i = (i >> 1))
	{
		if (threadIdx.x < i)
		{
			sh_sum[threadIdx.x] = sh_sum[threadIdx.x] + sh_sum[threadIdx.x + i];
		}
		__syncthreads();
	}

	if (threadIdx.x == 0)
	{
		devBuff[blockIdx.x] = sh_sum[0];
	}
}

__global__ void cuda_mean_merge(float * sumStep1, float * sumOut, int dataCount)
{
	int i = blockDim.x >> 1;

	extern __shared__ float sh_sum[];

	if (threadIdx.x < i)
	{
		int index1 = threadIdx.x + i;
		float secData = (index1 < dataCount) ? sumStep1[index1] : 0;
		sh_sum[threadIdx.x] = sumStep1[threadIdx.x] + secData;
	}
	__syncthreads();
	i = i >> 1;
	for (; i > 0; i = (i >> 1))
	{
		if (threadIdx.x < i)
		{
			sh_sum[threadIdx.x] = sh_sum[threadIdx.x] + sh_sum[threadIdx.x + i];
		}
		__syncthreads();
	}

	if (threadIdx.x == 0)
	{
		*sumOut = sh_sum[0];
	}

}

void cuda_mean_1d(uchar * devDataIn, float * devBuff, unsigned int length, float * hostMeanOut)
{
	int blockX = length > 1024 ? 1024 : length;
	dim3 blockSize(blockX, 1);

	int gridX = 1;
	if (length > 1024)
	{
		gridX = (length + 1023) / 1024;
	}

	//cout << "gridX = " << gridX << ",blockX = " << blockX<<endl;

	dim3 gridSize(gridX, 1);

	float * sumDevStep1 = devBuff;

	cuda_mean_kernel << <gridSize, blockSize, (blockX >> 1)*sizeof(float) >> >(devDataIn, sumDevStep1, length);

	int offsetBuff = ((gridX + 7) >> 3) << 3;

	float * meanDev = sumDevStep1 + offsetBuff;
	if (gridX > 1)
	{
		int blockX1 = nextPowerOfTwo(gridX);

		cuda_mean_merge << <1, blockX1, (blockX1 >> 1)*sizeof(float) >> >(sumDevStep1, meanDev, gridX);
	}

	cudaMemcpy(hostMeanOut, meanDev, sizeof(float), cudaMemcpyDeviceToHost);
	*hostMeanOut = (*hostMeanOut) / length;

#if DEBUG_MEAN
	float * step1Data = (float*)malloc(sizeof(float)*gridX);

	cudaMemcpy(step1Data, sumDevStep1, sizeof(float)*gridX, cudaMemcpyDeviceToHost);

	for (int i = 0; i < gridX; i++)
	{
		printf("%f, ", step1Data[i]);
	}
	printf("\n");
#endif
}
